import baostock as bs
import pandas as pd
import time
import logging
from tqdm.asyncio import tqdm_asyncio
import asyncio
from typing import List
# max concurrent process number
CONCURRENCY_FETCH = 2
CONCURRENCY_TRANSFORM = 2

# information columns needed
FIELDS = ("date,code,open,high,low,close," +
          "preclose,volume,amount,adjustflag,turn," +
          "tradestatus,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM,isST")

#date range of data
START_DATE = "2025-05-20"
END_DATE = "2025-06-01"

def get_stock_list():
    #return stock list of the component of csi500 index
    rs_com = bs.query_zz500_stocks()
    stock_list = []
    while (rs_com.error_code == '0') & rs_com.next():
        stock_list.append(rs_com.get_row_data()[1])

    return stock_list

async def get_data_response(code: str,semaphore):
     #return one stock data response
     async with semaphore:
         rs = bs.query_history_k_data_plus(
             code=code,
             fields=FIELDS,
             start_date=START_DATE,
             end_date=END_DATE,
             frequency="d", adjustflag="3")
     return rs

async def fetch_all_stocks(stock_list):
    # fetch all stock data
    semaphore = asyncio.Semaphore(CONCURRENCY_FETCH)  # 最大并发数=10
    tasks = [get_data_response(stock, semaphore) for stock in stock_list]
    results = await tqdm_asyncio.gather(*tasks)
    return results

async def transform_one_df(single_rs):
    data_list = []
    while (single_rs.error_code == '0') & single_rs.next():
        data_list.append(single_rs.get_row_data())
    result = pd.DataFrame(data_list, columns=single_rs.fields)
    return result

async def transform_df(rs_list):
    #transform the data into df
    tasks = [transform_one_df(single_rs) for single_rs in rs_list]
    result = await tqdm_asyncio.gather(*tasks)
    return result

async def main():
    lg = bs.login()

    stock_list = get_stock_list()[0:100]

    time1 = time.time()
    rs_list = await fetch_all_stocks(stock_list)
    time2 = time.time() - time1
    time1 = time.time()
    result_list = await transform_df(rs_list)
    time3 = time.time()-time1

    print(time2,time3)


if __name__ == "__main__":
    asyncio.run(main())